from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import os
import base64
from pydantic import BaseModel
from openai import OpenAI
import asyncio
import time

app = FastAPI()

from fastapi.middleware.cors import CORSMiddleware

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

api_key = "sk-*"  # 确保将其替换为您的实际 API 密钥
api_url = "https://api.manyi88.top/v1" #欢迎访问ManyiAPI聚合接口站：https://api.manyi88.top，注册优惠
client = OpenAI(api_key=api_key, base_url=api_url)

# Static folder
app.mount("/img", StaticFiles(directory="img"), name="img")

# Create img directory if it doesn't exist
os.makedirs("img", exist_ok=True)

@app.get("/", response_class=HTMLResponse)
async def read_root():
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Image Uploader</title>
        <style>
            body {
                font-family: Arial, sans-serif;
                display: flex;
                justify-content: center;
                align-items: center;
                height: 100vh;
                margin: 0;
                background-color: #f7f7f7;
            }
            .container {
                text-align: center;
                background: white;
                border-radius: 10px;
                box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
                padding: 20px;
                width: 90%;
                max-width: 800px;
            }
            textarea {
                width: 90%;
                border: 1px solid #ddd;
                border-radius: 5px;
                padding: 10px;
                margin: 10px 0;
                font-size: 14px;
                resize: none;
            }
            #preview {
                padding: 10px;
                margin: 10px;
            }
            button {
                padding: 10px 15px;
                border: none;
                border-radius: 5px;
                background-color: #007bff;
                color: white;
                font-size: 16px;
                cursor: pointer;
            }
            button:hover {
                background-color: #0056b3;
            }
            #loading {
                display: none;
                margin-top: 10px;
                padding: 10px;
            }
            #result {
                width: 90%;
                border: 1px solid #ddd;
                border-radius: 5px;
                padding: 10px;
                margin: auto;
                font-size: 14px;
                resize: none;
            }
           .file-upload {
            display: inline-block;
            padding: 10px 20px;
            background-color: #5daf1e;
            color: white;
            border: none;
            border-radius: 5px;
            cursor: pointer;
            }

            #fileInput {
                display: none; /* 隐藏原始文件输入控件 */
            }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>粘贴图片并添加描述</h1>
            <textarea id="imageData" rows="5" placeholder="粘贴图片的Base64数据..."></textarea>
            <br>
            <textarea id="prompt" rows="3" placeholder="输入提示信息..."></textarea>
            <button id="submitButton">提交</button>
            <label for="fileInput" class="file-upload">上传</label>
            <input type="file" id="fileInput" accept="image/*">
            <button id="clearButton">清空</button> <!-- Added clear button -->
            <button id="saveImageButton">保存图片</button>
            <select id="fileExtension">
                <option value="png">PNG</option>
                <option value="jpg">JPG</option>
                <option value="jpeg">JPEG</option>
                <option value="gif">GIF</option>
            </select>
            <div id="preview"></div>
            <div id="loading">提交中... <span class="loader"></span></div>
            <div id="result"></div>
        </div>

        <script>

            document.getElementById('fileInput').onchange = (event) => {
                const file = event.target.files[0];
                const reader = new FileReader();
                reader.onload = (event) => {
                    const imgData = event.target.result;
                    document.getElementById('imageData').value = imgData;
                    document.getElementById('preview').innerHTML = `<img src="${imgData}" style="max-width: 300px; max-height: 300px;" />`;
                };
                reader.readAsDataURL(file);
            };

            document.addEventListener('paste', (event) => {
                const items = (event.clipboardData || window.clipboardData).items;
                for (let i = 0; i < items.length; i++) {
                    if (items[i].type.indexOf("image") === 0) {
                        const file = items[i].getAsFile();
                        const reader = new FileReader();
                        reader.onload = (event) => {
                            const imgData = event.target.result;
                            document.getElementById('imageData').value = imgData;
                            document.getElementById('preview').innerHTML = `<img src="${imgData}" style="max-width: 300px; max-height: 300px;" />`;
                        };
                        reader.readAsDataURL(file);
                    }
                }
            });

            document.getElementById('submitButton').onclick = async () => {
                const imageData = document.getElementById('imageData').value;
                const promptInput = document.getElementById('prompt').value;

                const base64Data = imageData.split(',')[1];

                // Show loading indicator
                const loading = document.getElementById('loading');
                loading.style.display = 'block';

                const response = await fetch('/submit', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json'
                    },
                    body: JSON.stringify({
                        prompt: promptInput,
                        data: base64Data
                    })
                });

                const result = await response.json();
                loading.style.display = 'none'; // Hide loading indicator

                if (response.ok) {
                    document.getElementById('result').innerText = result.message;
                } else {
                    document.getElementById('result').innerText = '提交失败！';
                }
            };
            
            document.getElementById('clearButton').onclick = () => { // Added clear button functionality
            document.getElementById('imageData').value = '';
            document.getElementById('prompt').value = '';
            document.getElementById('preview').innerHTML = ''; // Clear preview
            document.getElementById('result').innerText = ''; // Clear result
            };

            document.getElementById('saveImageButton').onclick = async () => {
            const imageData = document.getElementById('imageData').value;
            const fileExtension = document.getElementById('fileExtension').value;

            if (!imageData) {
                alert('请先选择或粘贴图片');
                return;
            }

            const base64Data = imageData.split(',')[1];

            const response = await fetch('/save-image', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json'
                },
                body: JSON.stringify({
                    data: base64Data,
                    extension: fileExtension
                })
            });

            const result = await response.json();

            if (response.ok) {
                alert('图片已保存: ' + result.fileName);
            } else {
                alert('保存失败: ' + result.message);
            }
        };

        </script>
    </body>
    </html>
    """

class ImageUpload(BaseModel):
    prompt: str
    data: str

@app.post("/submit")
async def submit_image(image: ImageUpload):
    img_filename = "uploaded_image.png"
    file_path = os.path.join("img", img_filename)
    with open(file_path, "wb") as img_file:
        img_file.write(base64.b64decode(image.data))

    # Call OpenAI's GPT-4
    try:
        # print("here")
        response = client.chat.completions.create(
            model="gpt-4o-2024-08-06",
            # model="gpt-4o-mini-2024-07-18",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": image.prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image.data}"}}
                    ]
                }
            ],
            timeout=20
        )
        response_content = response.choices[0].message.content
        return {"message": response_content}
    except Exception as e:
        return {"message": f"Error in OpenAI API: {str(e)}"}

@app.post("/save-image")
async def save_image(data: dict):
    base64_data = data['data']
    extension = data['extension']
    timestamp = int(time.time())  # 使用时间戳
    img_filename = f"{timestamp}.{extension}"  # 使用时间戳作为文件名
    file_path = os.path.join("img", img_filename)

    try:
        with open(file_path, "wb") as img_file:
            img_file.write(base64.b64decode(base64_data))
        return {"fileName": img_filename}
    except Exception as e:
        return {"message": str(e)}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8008)

